import math
from collections import Counter

def euclidean_distance(x1, x2):
    """
    计算两个向量之间的欧几里得距离。
    """
    distance = 0.0
    for a, b in zip(x1, x2):
        distance += (a - b) ** 2
    return math.sqrt(distance)

class KNN:
    def __init__(self, k=3, task_type='classification'):
        """
        初始化 KNN 分类器/回归器。

        参数:
            k (int): 近邻数量。
            task_type (str): 任务类型, 'classification' (分类) 或 'regression' (回归)。
        """
        if k <= 0:
            raise ValueError("k 必须是正整数")
        self.k = k
        self.task_type = task_type.lower()
        if self.task_type not in ['classification', 'regression']:
            raise ValueError("task_type 必须是 'classification' 或 'regression'")
        
        # KNN 不训练模型，只存储数据
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):
        """
        拟合模型，即存储训练数据。

        参数:
            X (list of lists): 训练特征向量。
            y (list): 训练标签。
        """
        if len(X) != len(y):
            raise ValueError("特征向量 X 和标签 y 的长度必须相同")
        self.X_train = X
        self.y_train = y

    def predict(self, X):
        """
        对新样本进行预测。

        参数:
            X (list of lists): 待预测的特征向量。

        返回:
            list: 预测结果。
        """
        if self.X_train is None or self.y_train is None:
            raise ValueError("请先调用 fit 方法进行训练")
        
        predictions = []
        for x in X:
            # 1. 计算距离
            distances = [euclidean_distance(x, x_train) for x_train in self.X_train]
            
            # 2. 寻找 k 个最近邻的索引
            # argsort 返回的是数组值从小到大的索引值
            k_indices = sorted(range(len(distances)), key=lambda i: distances[i])[:self.k]
            
            # 3. 获取 k 个最近邻的标签
            k_nearest_labels = [self.y_train[i] for i in k_indices]

            # 4. 根据任务类型进行预测
            if self.task_type == 'classification':
                # 分类：投票表决
                most_common = Counter(k_nearest_labels).most_common(1)
                predictions.append(most_common[0][0])
            elif self.task_type == 'regression':
                # 回归：取平均值
                prediction = sum(k_nearest_labels) / len(k_nearest_labels)
                predictions.append(prediction)

        return predictions